#include "mincut_segmenter.h"

#include <iostream>
#include <imglib/bmpmorph.h>
#include <imglib/imageclass.h>
#include <imglib/pyramid.h>

namespace trm {

namespace mincut {

using std::clog;
using std::endl;
using std::flush;

void antialise(
	segmenter & s,
	segmenter::vmap_type const & vmap,
	params const & p,
	gray_mask & gmask)
{
	assert (p.gray_levels_exp <= 6);
	gmask.resize(vmap.dim());
	img::fill(gmask, 0);
	for (uint l = 0;; ++l) {
		s.inc_frame_gray_mask(vmap, gmask);
		if (l == p.gray_levels_exp)
			break;
		if (!p.silent)
			clog << "computing gray level " << (l+1) << "..." << endl;
		s.double_capacity_of_saturated_edges();
		s.solve();
	}

	uint m = (1 << (p.gray_levels_exp+2)) - 1;
	//table for converting values in gmask from [0, m) to [0, 256)
	std::vector<unsigned char> table(m);
	for (uint i = 0; i < m; ++i)
		table[i] = (255 * i) / (m - 1);
	for (gray_mask::iterator git = gmask.begin(); git != gmask.end(); ++git) {
		assert (*git < m);
		*git = table[*git];
	}
}

/// includes everything segment(...) should do except for pyramids
void segment_core(
	byte_image const & image,
	mask const & obligatory_foreground_pixels,
	mask const & obligatory_background_pixels,
	mask * foreground_only,
	mask * background_only,
	params const & p,
	gray_mask * gmask)
{
	if (!foreground_only && !background_only && !gmask)
		return;

	size_extent dim = image.dim();

	if (dim != obligatory_foreground_pixels.dim())
		throw std::runtime_error("dimensions of image and obligatory foreground pixels mask do not coincide");
	size_t ec1 = obligatory_foreground_pixels.set_bit_count();
	if (ec1 == 0)
		throw std::runtime_error("obligatory foreground pixels mask can't be empty");
	if (!p.silent)
		clog << "obligatory foreground pixels " << ec1 << endl;

	if (dim != obligatory_background_pixels.dim())
		throw std::runtime_error("dimensions of image and obligatory background pixels mask do not coincide");
	size_t ec0 = obligatory_background_pixels.set_bit_count();
	if (ec0 == 0)
		throw std::runtime_error("obligatory background pixels mask can't be empty");
	if (!p.silent)
		clog << "obligatory background pixels " << ec0 << endl;

	// input completed

	segmenter s(p.timing_tree, p.silent);

	segmenter::vmap_type vmap;

	if (!p.silent)
		clog << "building graph..." << endl;

	s.create_vmap(obligatory_foreground_pixels, obligatory_background_pixels, vmap);

	if (!p.silent)
		clog << "total vertices in graph: " << s.num_vertices() << endl;

	// create edges between pixels according to their similarity

	s.establish_links(
		image, image,
		vmap, vmap,
		p.hw, p.nn, p.aspect);

	if (!p.silent)
		clog << "total directional edges in graph: " << s.num_edges() << 
			"\ndone" << endl;

	s.solve();

	if (foreground_only)
		s.get_frame_mask(vmap, white_color, *foreground_only);

	if (background_only) 
		s.get_frame_mask(vmap, black_color, *background_only);

	if (gmask) 
		antialise(s, vmap, p, *gmask);
}

void segment(
	byte_image const & image,
	mask const & obligatory_foreground_pixels,
	mask const & obligatory_background_pixels,
	mask * foreground_only,
	mask * background_only,
	params const & p,
	gray_mask * gmask
)
{
	assert (p.max_motion >= 1);

	{
		mask m(obligatory_foreground_pixels);
		img::combine(m, obligatory_background_pixels, img::bit_and());
		if (m.set_bit_count() > 0)
			throw std::runtime_error("obligatory foreground and background pixels masks may not have set bits at the same locations");
	}

	if (p.pyramid_levels == 0)
		return segment_core(
			image, obligatory_foreground_pixels, obligatory_background_pixels,
			foreground_only, background_only,
			p, gmask);

	mswin::timing_mark downsampling_mark(p.timing_tree, "downsampling images & masks");
	img::pyramid<byte_image> image_pyramid(image, p.pyramid_levels, img::not_copy_original);
	img::pyramid<mask> obligatory_fmask_pyramid(obligatory_foreground_pixels, p.pyramid_levels, img::not_copy_original);
	img::pyramid<mask> obligatory_bmask_pyramid(obligatory_background_pixels, p.pyramid_levels, img::not_copy_original);
	for (uint l = p.pyramid_levels; l > 0; --l)
		img::combine(obligatory_bmask_pyramid.downsampled[l], obligatory_fmask_pyramid[l], img::bit_and_not());
	downsampling_mark.finish();

	mask fmask = obligatory_fmask_pyramid[p.pyramid_levels];
	mask bmask = obligatory_bmask_pyramid[p.pyramid_levels];

	for (uint l = p.pyramid_levels; l > 0; --l) {
		segment_core(image_pyramid[l], fmask, bmask, &fmask, &bmask, p, 0);

		size_extent updim = image_pyramid[l-1].dim();
		mask m;

		m.swap(fmask);
		fmask.resize(updim);
		fmask.fill(false);
		img::upsample(m, fmask, 2);
		img::erode4(fmask, p.max_motion);
		img::combine(fmask, obligatory_fmask_pyramid[l-1], img::bit_or());
		img::combine(fmask, obligatory_bmask_pyramid[l-1], img::bit_and_not()); //!new

		m.swap(bmask);
		bmask.resize(updim);
		bmask.fill(false);
		img::upsample(m, bmask, 2);
		img::erode4(bmask, p.max_motion);
		img::combine(bmask, obligatory_bmask_pyramid[l-1], img::bit_or());
		img::combine(bmask, obligatory_fmask_pyramid[l-1], img::bit_and_not()); //!new
	}

	segment_core(
		image, fmask, bmask,
		foreground_only, background_only,
		p,	gmask);
}

void adjust(
	byte_image const & image,
	mask const & initial_foreground_mask,
	mask * foreground_only,
	mask * background_only,
	params const & p,
	gray_mask * gmask)
{
	assert (p.max_motion >= 1);

	mask fmask, bmask; //obligatory masks on the current pyramid levels
	img::downsample(initial_foreground_mask, fmask, 1 << p.pyramid_levels);
	bmask = fmask;
	bmask.inverse();
	erode4(fmask, p.max_motion);
	erode4(bmask, p.max_motion);

	// the following code is much like to those in segment(...) function

	{
		mask m(bmask);
		img::combine(m, fmask, img::bit_and());
		if (m.set_bit_count() > 0)
			throw std::runtime_error("Either initial mask has too few foreground or background pixels, or maximum motion parameter is too large.");
	}

	if (p.pyramid_levels == 0)
		return segment_core(
			image, fmask, bmask,
			foreground_only, background_only,
			p, gmask);

	mswin::timing_mark downsampling_mark(p.timing_tree, "downsampling image");
	img::pyramid<byte_image> image_pyramid(image, p.pyramid_levels, img::not_copy_original);
	downsampling_mark.finish();

	for (uint l = p.pyramid_levels; l > 0; --l) {
		segment_core(image_pyramid[l], fmask, bmask, &fmask, &bmask, p, 0);

		size_extent updim = image_pyramid[l-1].dim();
		mask m;

		m.swap(fmask);
		fmask.resize(updim);
		fmask.fill(false);
		img::upsample(m, fmask, 2);
		img::erode4(fmask, p.max_motion);

		m.swap(bmask);
		bmask.resize(updim);
		bmask.fill(false);
		img::upsample(m, bmask, 2);
		img::erode4(bmask, p.max_motion);
	}

	segment_core(
		image, fmask, bmask,
		foreground_only, background_only,
		p,	gmask);
}

/// includes everything track(...) should do except for pyramids
void track_core(
	byte_image const & ref_image,
	mask const & ref_mask,
	byte_image const & image,
	mask const & obligatory_foreground_pixels,
	mask const & obligatory_background_pixels,
	mask * foreground_only,
	mask * background_only,
	params const & p,
	gray_mask * gmask)
{
	segmenter s(p.timing_tree, p.silent);

	segmenter::vmap_type ref_vmap, vmap;

	if (!p.silent)
		clog << "building graph..." << endl;

	s.create_vmap(ref_mask, ref_vmap);
	s.create_vmap(obligatory_foreground_pixels, obligatory_background_pixels, vmap);

	if (!p.silent)
		clog << "total vertices in graph: " << s.num_vertices() << endl;

	// create edges between pixels according to their similarity

	// inside target image
	s.establish_links(
		image, image,
		vmap, vmap,
		p.hw, p.nn, p.aspect);
	// "from" reference image "to" target image
	s.establish_links(
		ref_image, image,
		ref_vmap, vmap,
		p.interframe_hw, p.interframe_nn / 2, p.aspect
	);
	// "from" target image "to" reference image
	s.establish_links(
		image, ref_image,
		vmap, ref_vmap,
		p.interframe_hw, (p.interframe_nn + 1) / 2, p.aspect
	);

	if (!p.silent)
		clog << "total directional edges in graph: " << s.num_edges() << 
			"\ndone" << endl;

	s.solve();

	if (foreground_only)
		s.get_frame_mask(vmap, white_color, *foreground_only);

	if (background_only) 
		s.get_frame_mask(vmap, black_color, *background_only);

	if (gmask) 
		antialise(s, vmap, p, *gmask);
}

void track(
	byte_image const & ref_image,
	mask const & ref_mask,
	byte_image const & image,
	mask * foreground_only,
	mask * background_only,
	params const & p,
	gray_mask * gmask)
{
	assert (p.max_motion >= 1);
	assert (ref_image.dim() == ref_mask.dim());

	mask fmask, bmask; //obligatory masks on the current pyramid levels
	img::downsample(ref_mask, fmask, 1 << p.pyramid_levels);
	bmask = fmask;
	bmask.inverse();
	erode4(fmask, p.max_motion);
	erode4(bmask, p.max_motion);

	if (fmask.set_bit_count() == 0 || bmask.set_bit_count() == ref_mask.size())
		throw std::runtime_error("Either initial mask has too few foreground or background pixels, or maximum motion parameter is too large.");

	if (p.pyramid_levels == 0)
		return track_core(
			ref_image, ref_mask,
			image, fmask, bmask,
			foreground_only, background_only,
			p, gmask);
	assert (p.max_motion >= 1);

	mswin::timing_mark downsampling_mark(p.timing_tree, "downsampling images & masks");
	img::pyramid<byte_image> ref_image_pyramid(ref_image, p.pyramid_levels, img::not_copy_original);
	img::pyramid<mask> ref_mask_pyramid(ref_mask, p.pyramid_levels, img::not_copy_original);
	img::pyramid<byte_image> image_pyramid(image, p.pyramid_levels, img::not_copy_original);
	downsampling_mark.finish();

	for (uint l = p.pyramid_levels; l > 0; --l) {
		track_core(ref_image_pyramid[l], ref_mask_pyramid[l], 
			image_pyramid[l], fmask, bmask, &fmask, &bmask, p, 0);

		size_extent updim = image_pyramid[l-1].dim();
		mask m;

		m.swap(fmask);
		fmask.resize(updim);
		fmask.fill(false);
		img::upsample(m, fmask, 2);
		img::erode4(fmask, p.max_motion);

		m.swap(bmask);
		bmask.resize(updim);
		bmask.fill(false);
		img::upsample(m, bmask, 2);
		img::erode4(bmask, p.max_motion);
	}

	track_core(
		ref_image, ref_mask, image, fmask, bmask,
		foreground_only, background_only,
		p, gmask);
}

void run(
	byte_image const & ref_image,
	mask const & ref_mask,
	byte_image const & image,
	mask * foreground_only,
	mask * background_only,
	params const & p,
	gray_mask * gmask)
{
	if (p.source == source_both) {
		track(
			ref_image, ref_mask, image,
			foreground_only, background_only, p, gmask);
	}
	else {
		adjust(
			image, ref_mask,
			foreground_only, background_only, p, gmask);
	}
}

} //mincut

} //param
